import websocket
import hashlib
import base64
import hmac
import json
from urllib.parse import urlencode
import time
import ssl
from wsgiref.handlers import format_date_time
from datetime import datetime
from time import mktime
import _thread as thread
import pyaudio
import numpy as np
import queue
import threading
import sys
from threading import Event
import keyboard  # 新增：用于监听键盘事件


class XfyunSpeechRecognizer:
    """讯飞在线语音识别类封装 - 支持按s键开始录音，Enter键结束录音并识别"""

    def __init__(self, app_id, api_key, api_secret):
        # 初始化讯飞接口参数
        self.APPID = app_id
        self.APIKey = api_key
        self.APISecret = api_secret
        self.recording_results = ""  # 存储识别结果
        self.is_running = False  # 识别状态标识
        self.is_recording = False  # 录音状态标识

        # 音频帧状态常量
        self.STATUS_FIRST_FRAME = 0
        self.STATUS_CONTINUE_FRAME = 1
        self.STATUS_LAST_FRAME = 2

        # 初始化音频参数
        self.CHUNK = 1024  # 数据流块大小
        self.FORMAT = pyaudio.paInt16  # 16bit编码格式
        self.CHANNELS = 1  # 单声道
        self.RATE = 16000  # 采样频率

        # 音频缓冲区
        self.audio_buffer = queue.Queue()

        # 创建PyAudio实例
        self.p = pyaudio.PyAudio()

        # 控制事件
        self.start_event = Event()  # 开始录音事件
        self.stop_event = Event()   # 停止录音事件

        # 存储完整音频记录用于调试
        self.debug_mode = False
        self.debug_audio_frames = []
        self.callback = None  # 回调函数存储

        # 回调状态跟踪
        self.callback_in_progress = False
        self.callback_completed = False
        self.callback_lock = threading.Lock()
        self.ws_instance = None

    def set_callback(self, callback_func):
        """设置识别结果回调函数"""
        self.callback = callback_func

    def create_url(self):
        """生成websocket连接地址"""
        url = 'wss://ws-api.xfyun.cn/v2/iat'
        now = datetime.now()
        date = format_date_time(mktime(now.timetuple()))

        # 生成签名
        signature_origin = f"host: ws-api.xfyun.cn\n"
        signature_origin += f"date: {date}\n"
        signature_origin += "GET /v2/iat HTTP/1.1"

        signature_sha = hmac.new(
            self.APISecret.encode('utf-8'),
            signature_origin.encode('utf-8'),
            digestmod=hashlib.sha256
        ).digest()
        signature_sha = base64.b64encode(signature_sha).decode(encoding='utf-8')

        # 生成授权信息
        authorization_origin = (
            f"api_key=\"{self.APIKey}\", algorithm=\"hmac-sha256\", "
            f"headers=\"host date request-line\", signature=\"{signature_sha}\""
        )
        authorization = base64.b64encode(authorization_origin.encode('utf-8')).decode(encoding='utf-8')

        # 拼接URL参数
        v = {
            "authorization": authorization,
            "date": date,
            "host": "ws-api.xfyun.cn"
        }
        return f"{url}?{urlencode(v)}"

    def _audio_energy(self, data):
        """计算音频数据的能量"""
        try:
            audio_data = np.frombuffer(data, dtype=np.int16)
            squared = np.square(audio_data)
            mean = np.mean(squared) + 1e-10
            mean = max(0, mean)
            return np.sqrt(mean)
        except Exception as e:
            print(f"音频能量计算异常: {e}")
            return 0.0

    def _on_open(self, ws):
        """websocket连接建立后的回调"""

        def run(*args):
            print("---------------开始发送音频数据-----------------")
            status = self.STATUS_FIRST_FRAME

            # 从缓冲区获取音频数据
            audio_frames = []
            while not self.audio_buffer.empty():
                audio_frames.append(self.audio_buffer.get())

            try:
                # 发送音频数据
                for i, buf in enumerate(audio_frames):
                    if status == self.STATUS_FIRST_FRAME:
                        data = {
                            "common": {"app_id": self.APPID},
                            "business": {
                                "domain": "iat",
                                "language": "zh_cn",
                                "accent": "mandarin",
                                "vinfo": 1,
                                "vad_eos": 1000
                            },
                            "data": {
                                "status": 0,
                                "format": "audio/L16;rate=16000",
                                "audio": str(base64.b64encode(buf), 'utf-8'),
                                "encoding": "raw"
                            }
                        }
                        ws.send(json.dumps(data))
                        status = self.STATUS_CONTINUE_FRAME

                    elif status == self.STATUS_CONTINUE_FRAME:
                        data = {
                            "data": {
                                "status": 1,
                                "format": "audio/L16;rate=16000",
                                "audio": str(base64.b64encode(buf), 'utf-8'),
                                "encoding": "raw"
                            }
                        }
                        ws.send(json.dumps(data))

                    # 最后一帧
                    if i == len(audio_frames) - 1:
                        data = {
                            "data": {
                                "status": 2,
                                "format": "audio/L16;rate=16000",
                                "audio": str(base64.b64encode(buf), 'utf-8'),
                                "encoding": "raw"
                            }
                        }
                        ws.send(json.dumps(data))
                        break

            except Exception as e:
                print(f"发送音频异常: {e}")
            finally:
                print("---------------音频数据发送完毕-----------------")

        thread.start_new_thread(run, ())

    def _on_message(self, ws, message):
        """处理收到的websocket消息"""
        try:
            resp = json.loads(message)
            if resp.get("code") != 0:
                sid = resp.get("sid")
                err_msg = resp.get("message")
                print(f"sid:{sid} 错误: {err_msg} 错误码: {resp.get('code')}")
                return

            # 提取识别结果
            data = resp.get("data", {}).get("result", {}).get("ws", [])
            result = ""
            for item in data:
                for cw in item.get("cw", []):
                    result += cw.get("w", "")

            # 过滤无效结果
            if result not in ['', '。', '.。', ' .。', ' 。']:
                self.recording_results = result
                self.ws_instance = ws

                # 标记回调开始
                with self.callback_lock:
                    self.callback_in_progress = True
                    self.callback_completed = False

                # 调用回调函数
                if self.callback:
                    self.callback(result, self)

                # 等待回调完成
                while True:
                    with self.callback_lock:
                        if self.callback_completed:
                            break
                    time.sleep(0.01)

                # 回调完成后，主动关闭WebSocket连接
                if self.ws_instance:
                    self.ws_instance.close()

        except Exception as e:
            print(f"消息解析异常: {e}")

    def _on_error(self, ws, error):
        """处理错误信息"""
        print(f"WebSocket错误: {error}")
        self.is_running = False

    def _on_close(self, ws, close_status_code, close_msg):
        """处理连接关闭"""
        print("WebSocket连接已关闭")
        self.is_running = False
        self.ws_instance = None

        # 重置回调状态标志位
        with self.callback_lock:
            self.callback_in_progress = False
            self.callback_completed = False

        # 重置事件，准备下一次录音
        self.start_event.clear()
        self.stop_event.clear()
        
        print("\n准备就绪 - 按下 's' 键开始录音，按下 Enter 键结束录音并识别")

    def _recording_loop(self):
        """录音循环逻辑"""
        while True:
            # 等待开始录音信号
            print("等待 's' 键开始录音...")
            self.start_event.wait()
            self.start_event.clear()
            
            # 初始化录音状态
            self.is_recording = True
            self.audio_buffer.queue.clear()  # 清空缓冲区
            if self.debug_mode:
                self.debug_audio_frames = []
                
            print("开始录音... 按下 Enter 键结束录音")

            # 打开音频流
            stream = self.p.open(
                format=self.FORMAT,
                channels=self.CHANNELS,
                rate=self.RATE,
                input=True,
                frames_per_buffer=self.CHUNK
            )

            try:
                # 录音直到收到停止信号
                while self.is_recording and not self.stop_event.is_set():
                    buf = stream.read(self.CHUNK, exception_on_overflow=False)
                    self.audio_buffer.put(buf)
                    if self.debug_mode:
                        self.debug_audio_frames.append(buf)
                    time.sleep(0.001)  # 减少CPU占用

                # 停止录音
                print("停止录音，准备识别...")
                self.is_recording = False
                self.stop_event.clear()
                
                # 开始识别过程
                self._start_recognition_process()
                
            except Exception as e:
                print(f"录音异常: {e}")
            finally:
                stream.stop_stream()
                stream.close()

    def _start_recognition_process(self):
        """启动识别过程"""
        self.is_running = True
        websocket.enableTrace(False)
        ws_url = self.create_url()
        ws = websocket.WebSocketApp(
            ws_url,
            on_message=self._on_message,
            on_error=self._on_error,
            on_close=self._on_close
        )
        ws.on_open = self._on_open

        thread.start_new_thread(ws.run_forever, (), {"sslopt": {"cert_reqs": ssl.CERT_NONE}, "ping_timeout": 2})

    def start_listening(self):
        """开始监听按键和录音"""
        # 启动录音线程
        thread.start_new_thread(self._recording_loop, ())
        
        # 启动按键监听
        print("准备就绪 - 按下 's' 键开始录音，按下 Enter 键结束录音并识别")
        try:
            while True:
                # 监听's'键开始录音
                keyboard.wait('s')
                if not self.is_recording:
                    self.start_event.set()
                
                # 监听Enter键结束录音
                keyboard.wait('enter')
                if self.is_recording:
                    self.stop_event.set()
                
                time.sleep(0.1)
        except KeyboardInterrupt:
            self.stop_recognition()
            print("已手动停止程序")

    def get_result(self):
        """获取最新的识别结果"""
        return self.recording_results

    def stop_recognition(self):
        """停止录音识别"""
        self.is_recording = False
        self.is_running = False
        self.start_event.set()  # 唤醒可能阻塞的线程
        self.stop_event.set()
        print("已停止语音识别")

    def set_silence_threshold(self, threshold):
        """设置静音阈值（当前模式下未使用）"""
        self.SILENCE_THRESHOLD = threshold
        print(f"已设置静音阈值为: {threshold}")


# 使用示例
if __name__ == '__main__':
    # 替换为你的讯飞API信息
    APP_ID = "72b9a0ad"
    API_KEY = "59958ac9ffdaed40d02509f4bb593774"
    API_SECRET = "NmMyYTRlZjFhMmZkN2FkZjg0ODQ5NDVl"

    # 创建识别器实例
    recognizer = XfyunSpeechRecognizer(APP_ID, API_KEY, API_SECRET)

    # 可选：设置回调函数处理识别结果
    def handle_result(result, recognizer_instance):
        print(f"\n识别结果: {result}")
        # 标记回调完成
        with recognizer_instance.callback_lock:
            recognizer_instance.callback_completed = True

    recognizer.set_callback(handle_result)

    # 可选：调整参数
    # recognizer.set_silence_threshold(800)
    # recognizer.debug_mode = True

    try:
        recognizer.start_listening()
    except KeyboardInterrupt:
        recognizer.stop_recognition()
        print("程序已退出")